#include "vio/front_end/initialization.h"

namespace vio {
namespace initialization {

InitResult KltHomographyInit::AddFirstFrame(Frame::Ptr frame_ref) {
  Reset();
  DetectFeatures(frame_ref, px_ref_, f_ref_);
  if (px_ref_.size() < 100) {
    LOG(INFO) << "First image has less than 100 features, Retry in more "
                 "textured environment.";
    return FAILURE;
  }
  frame_ref_ = frame_ref;
  px_cur_.insert(px_cur_.begin(), px_ref_.begin(), px_ref_.end());
  return SUCCESS;
}

InitResult KltHomographyInit::AddSecondFrame(Frame::Ptr frame_cur) {
  TrackKLT(
      frame_ref_, frame_cur, px_ref_, px_cur_, f_ref_, f_cur_, disparities_);
  LOG(INFO) << "Init: KLT tracked " << disparities_.size() << " features";

  if (disparities_.size() < 50) {
    return FAILURE;
  }

  ComputeHomography(f_ref_,
                    f_cur_,
                    frame_ref_->cam_ptr_->ErrorMultiplier2(),
                    2.0,
                    inliers_,
                    xyz_in_cur_,
                    T_cur_from_ref_);
  LOG(INFO) << "Init: Homography RANSAC " << inliers_.size() << " inliers.";
  if (inliers_.size() < 40) {
    LOG(INFO) << "Init Warnning: " << 40 << " inliers minimum required.";
    return FAILURE;
  }

  // rescale the map such that the mean  scens depth is equal to the specified
  // scale
  std::vector<double> depth_vec;
  for (size_t i = 0; i < xyz_in_cur_.size(); i++) {
    depth_vec.push_back((xyz_in_cur_[i]).z());
    double scene_depth_median = math_utils::GetMedian(depth_vec);
    double scale = 1.0 / scene_depth_median;
    frame_cur->T_f_w_ = T_cur_from_ref_ * frame_ref_->T_f_w_;
    // rescale translation
    frame_cur->T_f_w_.translation() =
        -frame_cur->T_f_w_.rotationMatrix() *
        (frame_ref_->pos() + scale * (frame_cur->pos() - frame_ref_->pos()));
    // for each inlier create 3D point and add feature in both frames
    Sophus::SE3d T_world_cur = frame_cur->T_f_w_.inverse();
    for (std::vector<int>::iterator it = inliers_.begin(); it != inliers_.end();
         ++it) {
      Eigen::Vector2d px_cur(px_cur_[*it].x, px_cur_[*it].y);
      Eigen::Vector2d px_ref(px_ref_[*it].x, px_ref_[*it].y);
      if (frame_ref_->cam_ptr_->IsInFrame(px_cur.cast<int>(), 10) &&
          frame_ref_->cam_ptr_->IsInFrame(px_ref.cast<int>(), 10) &&
          xyz_in_cur_[*it].z() > 0) {
        Eigen::Vector3d pos = T_world_cur * (xyz_in_cur_[*it] * scale);
        Point::Ptr new_point = std::make_shared<Point>(pos);

        Feature::Ptr ftr_cur = std::make_shared<Feature>(
            frame_cur, new_point, px_cur, f_cur_[*it], 0);
        frame_cur->AddFeature(ftr_cur);
        new_point->AddFrameRef(ftr_cur);

        Feature::Ptr ftr_ref = std::make_shared<Feature>(
            frame_ref_, new_point, px_ref, f_ref_[*it], 0);
        frame_ref_->AddFeature(ftr_ref);
        new_point->AddFrameRef(ftr_ref);
      }
    }
  }
  return SUCCESS;
}

void KltHomographyInit::Reset() {
  px_cur_.clear();
  frame_ref_.reset();
}

void DetectFeatures(Frame::Ptr frame,
                    std::vector<cv::Point2f>& px_vec,
                    std::vector<Eigen::Vector3d>& f_vec) {
  Features new_features;
  feature_detection::FastDetector detector(
      frame->Img().cols, frame->Img().rows, 30, 3);
  detector.Detect(frame, frame->img_pyr_, 20.0, new_features);

  // now for all maximum corners, initialize a new seed
  px_vec.clear();
  px_vec.reserve(new_features.size());
  f_vec.clear();
  f_vec.reserve(new_features.size());
  // save keypoints and its' bearing
  std::for_each(
      new_features.begin(), new_features.end(), [&](Feature::Ptr ftr) {
        px_vec.push_back(cv::Point2f(ftr->px[0], ftr->px[1]));
        f_vec.push_back(ftr->bearing);
      });
}

void TrackKLT(Frame::Ptr frame_ref,
              Frame::Ptr frame_cur,
              std::vector<cv::Point2f>& px_ref,
              std::vector<cv::Point2f>& px_cur,
              std::vector<Eigen::Vector3d>& f_ref,
              std::vector<Eigen::Vector3d>& f_cur,
              std::vector<double>& disparities) {
  const double klt_win_size = 30.0;
  const int klt_max_iter = 30;
  const double klt_eps = 0.001;
  std::vector<uchar> status;
  std::vector<float> error;
  std::vector<float> min_eig_vec;
  // KLT terminal conditions
  cv::TermCriteria termcrit(
      cv::TermCriteria::COUNT + cv::TermCriteria::EPS, klt_max_iter, klt_eps);
  cv::calcOpticalFlowPyrLK(frame_ref->img_pyr_[0],
                           frame_cur->img_pyr_[0],
                           px_ref,
                           px_cur,
                           status,
                           error,
                           cv::Size2i(klt_win_size, klt_win_size),
                           4,
                           termcrit,
                           cv::OPTFLOW_USE_INITIAL_FLOW);
  std::vector<cv::Point2f>::iterator px_ref_it = px_ref.begin();
  std::vector<cv::Point2f>::iterator px_cur_it = px_cur.begin();
  std::vector<Eigen::Vector3d>::iterator f_ref_it = f_ref.begin();
  f_cur.clear();
  f_cur.reserve(px_cur.size());
  disparities.clear();
  disparities.reserve(px_cur.size());
  for (size_t i = 0; px_ref_it != px_ref.end(); ++i) {
    if (!status[i]) {
      px_ref_it = px_ref.erase(px_ref_it);
      px_cur_it = px_cur.erase(px_cur_it);
      f_ref_it = f_ref.erase(f_ref_it);
      continue;
    }
    f_cur.push_back(frame_cur->p2u(px_cur_it->x, px_cur_it->y));
    disparities.push_back(Eigen::Vector2d(px_ref_it->x - px_cur_it->x,
                                          px_ref_it->y - px_cur_it->y)
                              .norm());
    ++px_ref_it;
    ++px_cur_it;
    ++f_ref_it;
  }
}

void ComputeHomography(const std::vector<Eigen::Vector3d>& f_ref,
                       const std::vector<Eigen::Vector3d>& f_cur,
                       double focal_length,
                       double reprojection_threshold,
                       std::vector<int>& inliers,
                       std::vector<Eigen::Vector3d>& xyz_in_cur,
                       Sophus::SE3d& T_cur_from_ref) {
  std::vector<Eigen::Vector2d, Eigen::aligned_allocator<Eigen::Vector2d>>
      uv_ref(f_ref.size());
  std::vector<Eigen::Vector2d, Eigen::aligned_allocator<Eigen::Vector2d>>
      uv_cur(f_cur.size());
  for (size_t i = 0, i_max = f_ref.size(); i < i_max; ++i) {
    uv_ref[i] = math_utils::project2d(f_ref[i]);
    uv_cur[i] = math_utils::project2d(f_cur[i]);
  }
  vision_utils::Homography Homography(
      uv_ref, uv_cur, focal_length, reprojection_threshold);
  Homography.ComputeSE3FromMatches();
  std::vector<int> outliers;
  math_utils::ComputeInliers(f_cur,
                             f_ref,
                             Homography.T_c2_from_c1_.rotationMatrix(),
                             Homography.T_c2_from_c1_.translation(),
                             reprojection_threshold,
                             focal_length,
                             xyz_in_cur,
                             inliers,
                             outliers);
  T_cur_from_ref = Homography.T_c2_from_c1_;
}

}  // namespace initialization
}  // namespace vio